package dnl.persistence;

import au.com.bytecode.opencsv.CSVWriter;

import javax.sql.DataSource;
import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;

/**
 * @author Daniel Orr
 */
public class JdbcToCsv {

    private final String query;
    private final Connection connection;
    private String encoding = "UTF-8";
    private boolean writeColumnLabels;
    private DataConverter dataConverter = new DataConverter();

    private int flushInterval = 100;

    public JdbcToCsv(String query, Connection connection) {
        this.query = query;
        this.connection = connection;
    }

    public JdbcToCsv(String query, DataSource dataSource) {
        this.query = query;
        try {
            this.connection = dataSource.getConnection();
        } catch (SQLException e) {
            throw new IllegalStateException(e);
        }
    }

    public void setEncoding(String encoding) {
        this.encoding = encoding;
    }

    public void setWriteColumnLabels(boolean writeColumnLabels) {
        this.writeColumnLabels = writeColumnLabels;
    }

    public int getFlushInterval() {
        return flushInterval;
    }

    /**
     * Controls the flushing rate during the export. The underlying stream will be flushed according
     * to the number of rows read by the ResultSet. For example if {@code flushInterval} is 100, then
     * the data will be flushed every 100 rows. This helps controlling the memory consumption - the
     * larger the flush interval the bigger the memory consumption. Default value is 100.
     *
     * @param flushInterval number of result set rows between flushes.
     */
    public void setFlushInterval(int flushInterval) {
        this.flushInterval = flushInterval;
    }

    public void exportTo(File f) throws SQLException, IOException {
        try (CSVWriter csvWriter = new CSVWriter(new OutputStreamWriter(new FileOutputStream(f), encoding))) {
            final PreparedStatement ps = connection.prepareStatement(query);
            final ResultSet rs = ps.executeQuery();
            final ResultSetMetaData metaData = rs.getMetaData();
            final int columnCount = metaData.getColumnCount();
            if (writeColumnLabels) {
                writeColumnLabels(columnCount, metaData, csvWriter);
            }

            for (int rsCounter = 0; rs.next(); rsCounter++) {
                String[] rowData = new String[columnCount];
                for (int i = 0; i < columnCount; i++) {
                    final Object obj = rs.getObject(i + 1);
                    rowData[i] = dataConverter.fromDataObjectToString(obj);
                }
                if (rsCounter % flushInterval == 0) {
                    csvWriter.flush();
                }
                csvWriter.writeNext(rowData);
            }
            csvWriter.flush();
        }
    }

    private void writeColumnLabels(final int columnCount, final ResultSetMetaData metaData, final CSVWriter csvWriter) throws SQLException {
        String[] labels = new String[columnCount];
        for (int i = 0; i < columnCount; i++) {
            labels[i] = metaData.getColumnLabel(i+1);
        }
        csvWriter.writeNext(labels);
    }

}
